package slices

// GroupBy 将切片根据指定元素值进行分组
// T 切片类型
// K map的key类型
// arr 需要分组的切分
// fn 取对应key的函数
// map[K][]T 返回分组后的map
func GroupBy[T any, K comparable](arr []T, fn func(T) K) map[K][]T {
	res := make(map[K][]T)
	for _, e := range arr {
		k := fn(e)
		res[k] = append(res[k], e)
	}
	return res
}

// ToMap 将切片转换成map
// T 切片类型
// K map的key类型
// arr 需要转换的切分
// fn 取对应key的函数
// map[K]T 返回转换后的map
func ToMap[T any, K comparable](arr []T, fn func(T) K) map[K]T {
	res := make(map[K]T)
	for _, e := range arr {
		k := fn(e)
		res[k] = e
	}
	return res
}

// ToMapByKeyValue 将切片转换成map
// T 切片类型
// K map的key类型
// V map的value类型
// arr 需要转换的切分
// fn 取对应key的函数
// (K, V) 返回的map的key和value
// map[K]V 返回指定对应key和value的map
func ToMapByKeyValue[T any, K comparable, V any](arr []T, fn func(T) (K, V)) map[K]V {
	res := make(map[K]V)
	for _, e := range arr {
		k, v := fn(e)
		res[k] = v
	}
	return res
}

// Filter 从arr中筛选出符合条件的元素
// T 切片类型
// fn 取出符合的数据函数
// []T 返回符合的数据
func Filter[T any](arr []T, fn func(T) bool) []T {
	var n []T
	for _, e := range arr {
		if fn(e) {
			n = append(n, e)
		}
	}
	return n
}

// CountIf 统计数组中符合条件的个数
// arr 被判断数组
// fn 处理计数函数
// 返回：符合条件的计数值
func CountIf[T any](arr []T, fn func(T) bool) int {
	var count int
	for _, e := range arr {
		if fn(e) {
			count++
		}
	}
	return count
}

// Contains 判断
// arr 数组
// val 元素
// 返回 true：包含，false：不包含
func Contains[T comparable](arr []T, val T) bool {
	for _, v := range arr {
		if v == val {
			return true
		}
	}
	return false
}

// Map 根据函数处理数组并返回处理后的元素
func Map[T1, T2 any](arr []T1, fn func(T1) T2) []T2 {
	res := make([]T2, len(arr))
	for _, e := range arr {
		v := fn(e)
		res = append(res, v)
	}
	return res
}
